#include "ssabuilder.h"

#include <functional>
#include <string>
#include <unordered_map>
#include <vector>

#include "ast.h"
#include "irfunction.h"
#include "iroperand.h"
#include "utils.h"

void SSABuilder::Internalize(IRFunction* func) {
    std::unordered_set<std::string> activeGlobalVars;
    std::unordered_set<IRBlock*> vis;
    std::function<void(IRBlock*)> findCalls = [&](IRBlock* u) {
        if (vis.count(u)) return;
        vis.insert(u);
        for (auto ir : u->irList) {
            auto call = ir->As<IR::Call>();
            if (!call) continue;
            for (auto [name, _] : call->func->usedGlobalVars)
                activeGlobalVars.insert(name);
            for (auto [name, _] : call->func->globalEffects.var)
                activeGlobalVars.insert(name);
        }
        for (auto v : u->GetNext()) findCalls(v);
    };
    findCalls(func->entry);

    std::list<IR::IRBase*> load, store;
    std::unordered_map<std::string, IR::Variable> local;

    auto replace = [&](IR::Variable* var) {
        if (!var->isGlobal) return;
        if (activeGlobalVars.count(var->name)) return;
        if (!local.count(var->name)) {
            auto newVar = *var;
            newVar.name = "L" + var->name;
            newVar.isGlobal = false;

            auto irLoad = new IR::Operation();
            irLoad->op = ArithOp::Add;
            irLoad->rhs = new IR::Variable(*var);
            irLoad->dest = new IR::Variable(newVar);
            load.push_back(irLoad);

            auto irStore = new IR::Operation();
            irStore->op = ArithOp::Add;
            irStore->rhs = new IR::Variable(newVar);
            irStore->dest = new IR::Variable(*var);
            store.push_back(irStore);

            local[var->name] = newVar;
        }
        *var = local[var->name];
    };

    std::function<void(IRBlock*)> dfs = [&](IRBlock* u) {
        if (vis.count(u)) return;
        vis.insert(u);
        for (auto ir : u->irList) {
            for (auto v : ir->UsedVar()) replace(v);
            for (auto v : ir->ChangedVar()) replace(v);
        }
        for (auto v : u->GetNext()) dfs(v);
    };
    vis.clear();
    dfs(func->entry);

    func->entry->irList.insert(func->entry->irList.begin(), load.begin(),
                               load.end());
    func->exit->irList.insert(std::prev(func->exit->irList.end()),
                              store.begin(), store.end());
    func->SetParent();
}

void SSABuilder::InsertPhi(IRFunction* func) {
    auto DF = T.DF;

    std::unordered_map<std::string, IR::Variable*> varMap;
    std::unordered_map<std::string, std::vector<IR::IRBase*>> relIrs;
    for (auto block : T.nodes) {
        for (auto ir : block->irList) {
            for (auto var : ir->ChangedVar()) {
                if (var->isGlobal) continue;
                relIrs[var->name].push_back(ir);
                varMap[var->name] = var;
            }
        }
    }
    std::unordered_map<IRBlock*, std::string> vis;
    for (auto& [id, rel] : relIrs) {
        std::unordered_set<IRBlock*> phi;
        std::queue<IRBlock*> q;
        for (auto u : rel) {
            for (auto v : DF[u->parent]) q.push(v);
        }
        while (!q.empty()) {
            auto u = q.front();
            q.pop();
            phi.insert(u);
            if (vis[u] == id) continue;
            vis[u] = id;
            for (auto v : DF[u]) {
                q.push(v);
            }
        }
        for (auto u : phi) {
            auto* ir = new IR::Phi();
            ir->dest = new IR::Variable(*varMap[id]);
            for (auto v : u->GetPrev()) {
                ir->srcs[v] = new IR::Variable(*varMap[id]);
            }
            u->AddPhi(ir);
        }
    }
}

void SSABuilder::RenameVar(IRFunction* func) {
    std::unordered_map<std::string, int> count;
    std::unordered_map<std::string, std::vector<int>> history;
    std::unordered_map<IR::Variable*, int> res;
    auto top = [&](IR::Variable* var) {
        if (!history.count(var->name)) {
            history[var->name] = std::vector<int>{-1};
        }
        return history[var->name].back();
    };
    auto push = [&](IR::Variable* var, int i) {
        top(var);
        history[var->name].push_back(i);
    };
    auto pop = [&](IR::Variable* var) {
        int ret = history[var->name].back();
        history[var->name].pop_back();
        return ret;
    };

    auto renameUsed = [&](IR::Variable* var) {
        // Utils::Assert(!res.count(var), "variable has been renamed");
        res[var] = top(var);
    };
    auto renameChanged = [&](IR::Variable* var) {
        // Utils::Assert(!res.count(var), "variable has been renamed");
        int& version = count[var->name];
        res[var] = version;
        push(var, version);
        ++version;
    };
    std::function<void(IRBlock*)> rename = [&](IRBlock* block) {
        for (auto ir : block->phis) {
            for (auto var : ir->ChangedVar()) renameChanged(var);
        }
        for (auto ir : block->irList) {
            for (auto var : ir->UsedVar()) renameUsed(var);
            for (auto var : ir->ChangedVar()) renameChanged(var);
        }
        for (auto v : block->GetNext()) {
            for (auto phi : v->phis) {
                auto var = phi->srcs[block];
                renameUsed(var);
            }
        }

        for (auto v : T.dom[block]) rename(v);
        for (auto ir : block->irList) {
            for (auto var : ir->ChangedVar()) pop(var);
        }
        for (auto ir : block->phis) {
            for (auto var : ir->ChangedVar()) pop(var);
        }
    };
    for (auto var : func->args) renameChanged(var);
    rename(func->entry);

    for (auto [var, i] : res) {
        if (var->isGlobal) continue;
        var->name += "." + std::to_string(i);
    }
}

void SSABuilder::Build(IRFunction* func) {
    func->AnalEffects();
    Internalize(func);
    T.Build(func->entry);
    InsertPhi(func);
    RenameVar(func);
    func->Validate();
}

void* SSABuilder::Run(void* input) {
    auto module = (IRModule*)input;
    for (auto func : module->globalFuncs) {
        Build(func);
    }
    return module;
}
